import multiprocessing
import threading
from bisect import bisect
from collections import deque
from typing import Any, Callable, Iterable, Optional, Union

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import Button

from .repeating_timer import RepeatingTimer
from .scan_lines import LinesPoint, plot_summary
from .utils import convert_to_iterator

LinesBufferData = Optional[tuple[np.ndarray, np.ndarray]]
LinesViewData = Optional[tuple[np.ndarray, np.ndarray, Optional[tuple[np.ndarray, ...]]]]


class ScanLinesBuffer:
    def __init__(self, size: Optional[int] = 1024):
        self._lock = threading.RLock()
        self._condition = threading.Condition(self._lock)

        self._version = 0
        self._points = deque[LinesPoint](maxlen=size)

    def append(self, point: LinesPoint):
        with self._lock:
            self._version += 1
            self._points.append(point)
            self._condition.notify_all()

    def clear(self):
        with self._lock:
            self._version += 1
            self._points.clear()
            self._condition.notify_all()

    def read(self) -> tuple[int, LinesBufferData]:
        with self._lock:
            version = self._version
            buffer = tuple(self._points)

        if len(buffer) == 0:
            return version, None

        x_logs, ys_logs = zip(*buffer)
        x_logs = np.asarray(x_logs, dtype=np.float32)  # [len]
        ys_logs = np.asarray(ys_logs, dtype=np.float32)  # [len, (ys_n)]
        return version, (x_logs, ys_logs)

    def wait_and_read(self, version: int) -> tuple[int, LinesBufferData]:
        with self._lock:
            while version == self._version:
                self._condition.wait()
            return self.read()


def convert_to_view_data(
    data: LinesBufferData, *,
    win: Optional[float] = None,
    bins: Optional[int] = None,
) -> LinesViewData:
    if data is None:
        return None
    x_logs, ys_logs = data

    if np.ndim(ys_logs) == 1:
        ys_logs = np.expand_dims(ys_logs, axis=-1)
    ys_n = np.shape(ys_logs)[-1]

    # histograms
    if bins is not None:
        histograms = [np.histogram(ys_logs[:, yi], bins=bins) for yi in range(ys_n)]
    else:
        histograms = None

    # curves
    if win and len(x_logs) > 0:
        head = bisect(x_logs, x_logs[-1] - win)
        x_logs = x_logs[head:]
        ys_logs = ys_logs[head:]

    return x_logs, ys_logs, histograms


class ScanLinesView:
    def __init__(self,
        figure: plt.Figure, *,
        win: float,
        hist: bool = False,
        title: Optional[str] = None,
    ):
        self._figure = figure
        self._win = win
        self._hist = hist

        self._init_figure(figure, hist)
        self._axes_curves, self._axes_histograms = self._init_axes(figure, hist, title)
        self._play_button, self._clear_button = self._init_buttons(figure)
        self._play_button.on_clicked(lambda _: self._on_play_clicked())
        self._clear_button.on_clicked(lambda _: self._on_clear_clicked())

        self._ys_lines: Optional[list[plt.Line2D]] = None
        self._ys_bars: Optional[list[tuple[plt.Rectangle, ...]]] = None

        self._paused = True
        self.request_pause: Optional[Callable[[], Any]] = None
        self.request_resume: Optional[Callable[[], Any]] = None
        self.request_clear: Optional[Callable[[], Any]] = None

    @staticmethod
    def _init_figure(figure: plt.Figure, hist: bool):
        if hist:
            figure.set_size_inches(9, 4)
        else:
            figure.set_size_inches(8, 4)
        return figure

    @staticmethod
    def _init_axes(figure: plt.Figure, hist: bool, title: Optional[str] = None) -> tuple[plt.Axes, Optional[plt.Axes]]:
        if hist:
            axes_curves, axes_histograms = figure.subplots(1, 2,
                sharey='row', gridspec_kw=dict(width_ratios=[3, 1], wspace=0.05))
            axes_curves.grid()
            axes_histograms.grid()
        else:
            axes_curves = figure.gca()
            axes_curves.grid()
            axes_histograms = None

        if title is not None:
            axes_curves.set_title(title)

        return axes_curves, axes_histograms

    @staticmethod
    def _init_buttons(figure: plt.Figure) -> tuple[plt.Button, plt.Button]:
        figure.subplots_adjust(0.1, 0.1, 0.95, 0.85)
        play_button = Button(figure.add_axes([0.01, 0.9, 0.09, 0.09]), "resume")
        clear_button = Button(figure.add_axes([0.11, 0.9, 0.09, 0.09]), "clear")
        return play_button, clear_button

    @staticmethod
    def _init_lines(axes: plt.Axes, ys_n: int) -> list[plt.Line2D]:
        return [axes.plot(np.zeros([0]), zorder=2 + (ys_n - yi) / ys_n)[0]
                for yi in range(ys_n)]

    @staticmethod
    def _init_bars(axes: plt.Axes, ys_n: int, bars_n: int) -> list[tuple[plt.Rectangle, ...]]:
        return [axes.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def update(self, data: LinesViewData):
        if data is None:  # clear
            self._axes_curves.cla()
            self._axes_curves.grid()
            self._ys_lines = None
            if self._axes_histograms:
                self._axes_histograms.cla()
                self._axes_histograms.grid()
                self._ys_bars = None
            self._figure.canvas.draw_idle()
            return

        x_curve, ys_curve, ys_histograms = data
        if self._ys_lines is None:
            ys_n = np.shape(ys_curve)[-1]
            if ys_n == 0:
                return

            if self._hist:
                bars_n = len(ys_histograms[0][0])
                if bars_n == 0:
                    return
            else:
                bars_n = None

            self._ys_lines = self._init_lines(self._axes_curves, ys_n)
            if self._axes_histograms:
                self._ys_bars = self._init_bars(self._axes_histograms, ys_n, bars_n)

        # update lines
        for yi, line in enumerate(self._ys_lines):
            line.set_data(x_curve, ys_curve[:, yi])

        # adjust lines x range
        min_x, max_x = np.min(x_curve), np.max(x_curve)
        if max_x < min_x + self._win:
            max_x = min_x + self._win
        else:
            min_x = max_x - self._win
        self._axes_curves.set_xlim(min_x, max_x)

        # update histogram
        if self._axes_histograms:
            max_height = 0
            for yi, (y_bars, (heights, edges)) in enumerate(zip(self._ys_bars, ys_histograms)):
                for yi, bar in enumerate(y_bars):
                    bar.set_y(edges[yi])
                    bar.set_height(edges[yi + 1] - edges[yi])
                    bar.set_width(heights[yi])
                max_height = max(max_height, np.max(heights))

            # adjust histogram x range
            self._axes_histograms.set_xlim(0, max_height * 1.1)

        # adjust both y range
        min_y, max_y = np.min(ys_curve), np.max(ys_curve)
        if max_y > min_y:
            padding_y = (max_y - min_y) / 32
            self._axes_curves.set_ylim(min_y - padding_y, max_y + padding_y)
            if self._hist:
                self._axes_histograms.set_ylim(min_y - padding_y, max_y + padding_y)

        # redraw
        self._figure.canvas.draw_idle()

    def pause(self):
        if self._paused:
            return
        self._paused = True
        self._play_button.label.set_text("resume")

    def resume(self):
        if not self._paused:
            return
        self._paused = False
        self._play_button.label.set_text("pause")

    def _on_play_clicked(self):
        if self._paused:
            if self.request_resume:
                self.request_resume()
        else:
            if self.request_pause:
                self.request_pause()

    def _on_clear_clicked(self):
        if self.request_clear:
            self.request_clear()


class ScanLinesPopup:
    def __init__(self,
        producer: Union[Iterable[LinesPoint], Callable[[], LinesPoint]], *,
        produce_interval: float = 1 / 128,
        frame_interval: float = 1 / 32,
        win: float = 200,
        bins: Optional[int] = None,
        title: Optional[str] = None,
        start: bool = True,
        buffer_size: int = 1024,
    ):
        """ Scan the data from `producer` as lines in a popup window

        :param producer: an iterator (generator) or callable that produces data points in form (x, ys)
        :param produce_interval: the interval between calls of `producer`
        :param frame_interval: the interval between gui updates
        :param win: the window size of the scan
        :param bins: the number of bins for the histogram. if None, no histogram is shown
        :param title: the title of the plot
        :param start: whether to start immediately
        :param buffer_size: the number of buffered data points. if None, the buffer is unbounded. 
        """
        self._iterator = convert_to_iterator(producer)
        self._data_connection, sub_data_connection = multiprocessing.Pipe()
        self._event_connection, sub_event_connection = multiprocessing.Pipe()
        self._subprocess = multiprocessing.Process(
            target=self.subprocess_func,
            args=(sub_data_connection, sub_event_connection,
                  frame_interval, win, bins, title, buffer_size))
        self._subprocess.start()
        self._buffer = ScanLinesBuffer(buffer_size)

        def tick_func():
            point = next(self._iterator)
            try:
                self._data_connection.send(point)
            except (OSError, BrokenPipeError):
                self._timer.stop()
                return
            self._buffer.append(point)

        self._timer = RepeatingTimer(produce_interval, tick_func)

        def event_thread_func():
            while True:
                try:
                    cmd = self._event_connection.recv()
                except EOFError:
                    break
                if cmd == "pause":
                    self.pause()
                elif cmd == "resume":
                    self.resume()
                elif cmd == "clear":
                    self.clear()

        self._event_thread = threading.Thread(target=event_thread_func)
        self._event_thread.start()

        self._win = win
        self._bins = bins
        self._title = title

        if start:
            self.resume()

    @staticmethod
    def subprocess_func(
        data_connection,
        event_connection,
        frame_interval: float,
        win: float,
        bins: Optional[int],
        title: Optional[str],
        buffer_size: int,
    ):
        matplotlib.use('QtAgg')

        buffer = ScanLinesBuffer(buffer_size)

        def data_thread_func():
            while True:
                point = data_connection.recv()
                buffer.append(point)

        data_thread = threading.Thread(target=data_thread_func, daemon=True)
        data_thread.start()

        figure = plt.figure()
        view = ScanLinesView(figure, win=win, hist=bins is not None, title=title)
        view_version = -1

        def tick_func():
            nonlocal view_version
            view_version, data = buffer.wait_and_read(view_version)
            view_data = convert_to_view_data(data, win=win, bins=bins)
            view.update(view_data)

        timer = RepeatingTimer(frame_interval, tick_func)
        timer.start()

        view.request_pause = lambda: event_connection.send("pause")
        view.request_resume = lambda: event_connection.send("resume")
        view.request_clear = lambda: event_connection.send("clear")

        def event_thread_func():
            while True:
                cmd = event_connection.recv()
                if cmd == "pause":
                    view.pause()
                elif cmd == "resume":
                    view.resume()
                elif cmd == "clear":
                    buffer.clear()

        event_thread = threading.Thread(target=event_thread_func, daemon=True)
        event_thread.start()

        plt.show()

    def resume(self):
        self._timer.start()
        self._event_connection.send("resume")

    def pause(self):
        self._timer.stop()
        self._event_connection.send("pause")

    def clear(self):
        self._buffer.clear()
        self._event_connection.send("clear")

    def wait(self):
        if self._subprocess.is_alive():
            self._subprocess.join()
        self._data_connection.close()
        self._event_connection.close()
        self._timer.join()
        self._event_thread.join()

    def close(self):
        if self._subprocess.is_alive():
            self._subprocess.terminate()
            self._subprocess.join()
        self._data_connection.close()
        self._event_connection.close()
        self._timer.join()
        self._event_thread.join()

    def plot_summary(self):
        _, data = self._buffer.read()
        view_data = convert_to_view_data(data, win=self._win, bins=self._bins)
        plot_summary(view_data, win=self._win, hist=self._bins is not None, title=self._title)

    def wait_and_plot_summary(self):
        self.wait()
        self.plot_summary()

    def close_and_plot_summary(self):
        self.close()
        self.plot_summary()



def run_scan_with_popup(
        scanranges,                     # e.g. np.linspace(...)
        update_function,                # your slow, blocking instrument call
        *,
        title="Scan",                   # title on the live & summary plots
        gui_fps=30                      # how fast the GUI repaints
    ):
    """
    Launch a ScanLinesPopup that closes itself when the scan is done.
    """
    # ────────────────────────────────────────────────────────────────────
    # 1) Build the popup first, with an *empty* producer and start=False
    #    (we will feed data manually, so no need for the built‑in timer).
    # ────────────────────────────────────────────────────────────────────
    popup = ScanLinesPopup(
        producer=[],                 # placeholder, we'll push points ourselves
        produce_interval=1e9,        # irrelevant because we never start its timer
        frame_interval=1.0/gui_fps,  # GUI refresh rate (child process)
        win=float(scanranges[-1] - scanranges[0]),
        bins=None,                   # set an int for histograms if you like
        title=title,
        start=False,                 # don't start the (unused) producer timer
        buffer_size=len(scanranges), # keep all points for the summary
    )

    # ────────────────────────────────────────────────────────────────────
    # 2) Define a *worker thread* that actually performs the scan.
    #    It pushes each LinesPoint into the popup's IPC pipe
    #    and appends to the main‑process buffer for the summary.
    # ────────────────────────────────────────────────────────────────────
    def scan_worker():
        try:
            for x in scanranges:
                ys = np.atleast_1d(update_function(x)).astype(np.float32)

                # make the data point
                point: LinesPoint = (float(x), ys)

                # send to child process (live graph)
                popup._data_connection.send(point)

                # store locally (for summary plot later)
                popup._buffer.append(point)

            # all points done → auto‑close window *and* draw summary
            popup.close_and_plot_summary()

        except Exception as err:
            # any instrument error: print and tear everything down
            print("Scan aborted:", err)
            popup.close()

    # ────────────────────────────────────────────────────────────────────
    # 3) Kick off the worker *after* the GUI process is alive.
    #    (The child process was started inside ScanLinesPopup.__init__.)
    # ────────────────────────────────────────────────────────────────────
    threading.Thread(target=scan_worker, daemon=True).start()
